{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import numpy.linalg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib widget\n",
    "np.set_printoptions(precision=8, suppress=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import data_receiver\n",
    "from mathlib import *\n",
    "from plotlib import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# sampling rate\n",
    "dt = 0.01    # s\n",
    "\n",
    "# the initialization interval\n",
    "ts = 1    # s"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Pull Data From Phone\n",
    "data order: gyroscorpe, accelerometer, magnetometer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "r = data_receiver.Receiver()\n",
    "\n",
    "data = []\n",
    "\n",
    "for line in r.receive():\n",
    "    data.append(line.split(','))\n",
    "\n",
    "data = np.array(data, dtype = np.float)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Initialization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# discard the first and last few readings\n",
    "# for some reason they fluctuate a lot\n",
    "w = data[10:-10, 0:3]\n",
    "a = data[10:-10, 3:6]\n",
    "m = data[10:-10, 6:9]\n",
    "\n",
    "if(np.shape(w)[0] < ts/dt):\n",
    "    print(\"not enough data for intialization!\")\n",
    "\n",
    "# gravity\n",
    "gn = a[:int(ts/dt)].mean(axis = 0)\n",
    "gn = -gn[:, np.newaxis]\n",
    "g0 = np.linalg.norm(gn)  # save the initial magnitude of gravity\n",
    "\n",
    "# magnetic field\n",
    "mn = m[:int(ts/dt)].mean(axis = 0)\n",
    "mn = Normalized(mn)[:, np.newaxis]  # magnitude is not important\n",
    "\n",
    "avar = a[:int(ts/dt)].var(axis=0)\n",
    "wvar = w[:int(ts/dt)].var(axis=0)\n",
    "mvar = m[:int(ts/dt)].var(axis=0)\n",
    "print('acc var: ', avar, ', ', np.linalg.norm(avar))\n",
    "print('ang var: ', wvar, ', ', np.linalg.norm(wvar))\n",
    "print('mag var: ', mvar, ', ', np.linalg.norm(mvar))\n",
    "\n",
    "# cut the initialization data\n",
    "w = w[int(ts/dt) - 1:] - w[:int(ts/dt)].mean(axis=0)\n",
    "a = a[int(ts/dt):]\n",
    "m = m[int(ts/dt):]\n",
    "\n",
    "sample_number = np.shape(a)[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a_filtered, w_filtered = Filt_signal((a, w), dt=dt, wn=10, btype='lowpass')\n",
    "plot_signal([a, a_filtered], [w, w_filtered], [m])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Kalman Filter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gyro_noise = 10 * np.linalg.norm(wvar)\n",
    "acc_noise = 10 * np.linalg.norm(avar)\n",
    "mag_noise = 10 * np.linalg.norm(mvar)\n",
    "\n",
    "P = 1e-10 * I(4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "outputPrepend"
    ]
   },
   "outputs": [],
   "source": [
    "a_nav = []\n",
    "orientations = []\n",
    "\n",
    "q = np.array([[1., 0., 0., 0.]]).T\n",
    "orin = -gn / np.linalg.norm(gn)\n",
    "\n",
    "t = 0\n",
    "while t < sample_number:\n",
    "    wt = w[t, np.newaxis].T\n",
    "    at = a[t, np.newaxis].T\n",
    "    mt = m[t, np.newaxis].T \n",
    "    mt = Normalized(mt)\n",
    "\n",
    "    # Propagation\n",
    "    Ft = F(q, wt, dt)\n",
    "    Gt = G(q)\n",
    "    Q = (gyro_noise * dt)**2 * Gt @ Gt.T\n",
    "    \n",
    "    q = Ft @ q\n",
    "    q = Normalized(q)\n",
    "    P = Ft @ P @ Ft.T + Q    \n",
    "\n",
    "    # Measurement Update\n",
    "    # Use only normalized measurements to reduce error!\n",
    "    \n",
    "    # acc and mag prediction\n",
    "    pa = Normalized(-Rotate(q) @ gn)\n",
    "    pm = Normalized(Rotate(q) @ mn)\n",
    "\n",
    "    # Residual\n",
    "    Eps = np.vstack((Normalized(at), mt)) - np.vstack((pa, pm))\n",
    "    \n",
    "    # internal error + external error\n",
    "    Ra = [(acc_noise / np.linalg.norm(at))**2 + (1 - g0 / np.linalg.norm(at))**2] * 3\n",
    "    Rm = [mag_noise**2] * 3\n",
    "    R = np.diag(Ra + Rm)\n",
    "    \n",
    "    Ht = H(q, gn, mn)\n",
    "\n",
    "    S = Ht @ P @ Ht.T + R\n",
    "    K = P @ Ht.T @ np.linalg.inv(S)\n",
    "    q = q + K @ Eps\n",
    "    P = P - K @ Ht @ P\n",
    "    \n",
    "    # Post Correction\n",
    "    q = Normalized(q)\n",
    "    P = 0.5 * (P + P.T)  # make sure P is symmertical\n",
    "    \n",
    "    conj = -I(4)\n",
    "    conj[0, 0] = 1\n",
    "    an = Rotate(conj @ q) @ at + gn\n",
    "    ori = Rotate(conj @ q) @ orin\n",
    "\n",
    "    a_nav.append(an.T[0])\n",
    "    orientations.append(ori.T[0])\n",
    "\n",
    "    t += 1\n",
    "\n",
    "a_nav = np.array(a_nav)\n",
    "orientations = np.array(orientations)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Accelerometer Bias/Error Correction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a_threshold = 0.2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "outputPrepend"
    ]
   },
   "outputs": [],
   "source": [
    "t_start = 0\n",
    "for t in range(sample_number):\n",
    "    at = a_nav[t]\n",
    "    if np.linalg.norm(at) > a_threshold:\n",
    "        t_start = t\n",
    "        break\n",
    "\n",
    "t_end = 0\n",
    "for t in range(sample_number - 1, -1,-1):\n",
    "    at = a_nav[t]\n",
    "    if np.linalg.norm(at - a_nav[-1]) > a_threshold:\n",
    "        t_end = t\n",
    "        break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('motion starts at: ', t_start)\n",
    "print('motion ends at: ', t_end)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "an_drift = a_nav[t_end:].mean(axis=0)\n",
    "an_drift_rate = an_drift / (t_end - t_start)\n",
    "\n",
    "for i in range(t_end - t_start):\n",
    "    a_nav[t_start + i] -= (i+1) * an_drift_rate\n",
    "\n",
    "for i in range(sample_number - t_end):\n",
    "    a_nav[t_end + i] -= an_drift"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "filtered_a_nav, = Filt_signal([a_nav], dt=dt, wn=(0.01, 15), btype='bandpass')\n",
    "plot_3([a_nav, filtered_a_nav])\n",
    "# plot_3([a_nav])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Zero Velocity Update"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "outputPrepend"
    ]
   },
   "outputs": [],
   "source": [
    "velocities = []\n",
    "prevt = -1\n",
    "still_phase = False\n",
    "\n",
    "v = np.zeros((3, 1))\n",
    "t = 0\n",
    "while t < sample_number:\n",
    "    at = filtered_a_nav[t, np.newaxis].T\n",
    "\n",
    "    if np.linalg.norm(at) < a_threshold:\n",
    "        if not still_phase:\n",
    "            predict_v = v + at * dt\n",
    "\n",
    "            v_drift_rate = predict_v / (t - prevt)\n",
    "            for i in range(t - prevt - 1):\n",
    "                velocities[prevt + 1 + i] -= (i + 1) * v_drift_rate.T[0]\n",
    "\n",
    "        v = np.zeros((3, 1))\n",
    "        prevt = t\n",
    "        still_phase = True\n",
    "    else:\n",
    "        v = v + at * dt\n",
    "        still_phase = False\n",
    "    \n",
    "    t += 1\n",
    "    velocities.append(v.T[0])\n",
    "velocities = np.array(velocities)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_3([velocities])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Integration To Get Position"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "positions = []\n",
    "p = np.array([[0, 0, 0]]).T\n",
    "\n",
    "t = 0\n",
    "while t < sample_number:\n",
    "    at = filtered_a_nav[t, np.newaxis].T\n",
    "    vt = velocities[t, np.newaxis].T\n",
    "\n",
    "    p = p + vt * dt + 0.5 * at * dt**2\n",
    "    positions.append(p.T[0])\n",
    "\n",
    "    t += 1\n",
    "\n",
    "positions = np.array(positions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_3D([[positions, 'position']])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_3([positions])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Close All Graphs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.close('all')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.7-final"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}